import pygame
import random
import sys

# 定义常量
SCREEN_WIDTH = 640
SCREEN_HEIGHT = 480
BLOCK_SIZE = 20
FPS = 10

# 定义颜色常量
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
RED = (255, 0, 0)
GREEN = (0, 255, 0)

# 初始化 pygame
pygame.init()

# 设置游戏窗口
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("贪吃蛇")

# 设置字体
font = pygame.font.SysFont(None, 30)

# 定义贪吃蛇类
class Snake:
    def __init__(self):
        self.body = [pygame.Rect(200, 200, BLOCK_SIZE, BLOCK_SIZE),
                     pygame.Rect(220, 200, BLOCK_SIZE, BLOCK_SIZE),
                     pygame.Rect(240, 200, BLOCK_SIZE, BLOCK_SIZE)]
        self.color = GREEN

    def draw(self, screen):
        for block in self.body:
            pygame.draw.rect(screen, self.color, block)

    def move(self, direction):
        if direction == "up":
            self.body.insert(0, pygame.Rect(self.body[0].left, self.body[0].top-BLOCK_SIZE, BLOCK_SIZE, BLOCK_SIZE))
        elif direction == "down":
            self.body.insert(0, pygame.Rect(self.body[0].left, self.body[0].top+BLOCK_SIZE, BLOCK_SIZE, BLOCK_SIZE))
        elif direction == "left":
            self.body.insert(0, pygame.Rect(self.body[0].left-BLOCK_SIZE, self.body[0].top, BLOCK_SIZE, BLOCK_SIZE))
        elif direction == "right":
            self.body.insert(0, pygame.Rect(self.body[0].left+BLOCK_SIZE, self.body[0].top, BLOCK_SIZE, BLOCK_SIZE))
        self.body.pop()

    def grow(self):
        self.body.append(pygame.Rect(self.body[-1].left, self.body[-1].top, BLOCK_SIZE, BLOCK_SIZE))

    def hit_wall(self):
        if self.body[0].left < 0 or self.body[0].right > SCREEN_WIDTH or \
            self.body[0].top < 0 or self.body[0].bottom > SCREEN_HEIGHT:
            return True
        else:
            return False

    def hit_self(self):
        for block in self.body[1:]:
            if self.body[0].colliderect(block):
                return True
        else:
            return False

# 定义食物类
class Food:
    def __init__(self):
        self.rect = pygame.Rect(0, 0, BLOCK_SIZE, BLOCK_SIZE)
        self.color = RED
        self.spawn()

    def draw(self, screen):
        pygame.draw.rect(screen, self.color, self.rect)

    def spawn(self):
        x = random.randint(0, SCREEN_WIDTH//BLOCK_SIZE-1)
        y = random.randint(0, SCREEN_HEIGHT//BLOCK_SIZE-1)
        self.rect.topleft = (x*BLOCK_SIZE, y*BLOCK_SIZE)

# 定义得分函数
def draw_score(screen, score):
    score_text = font.render("Score: " + str(score), True, WHITE)
    screen.blit(score_text, (10, 10))

# 定义文字绘制函数
def draw_text(screen, text, x, y):
    text_obj = font.render(text, True, WHITE)
    screen.blit(text_obj, (x, y))

def main():
    # 初始化贪吃蛇和食物
    snake = Snake()
    food = Food()
    # 定义变量
    direction = "up"
    score = 0
    game_over = False

    # 设置游戏主循环
    clock = pygame.time.Clock()
    while not game_over:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                game_over = True
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_UP and direction != "down":
                    direction = "up"
                elif event.key == pygame.K_DOWN and direction != "up":
                    direction = "down"
                elif event.key == pygame.K_LEFT and direction != "right":
                    direction = "left"
                elif event.key == pygame.K_RIGHT and direction != "left":
                    direction = "right"

        # 移动贪吃蛇
        snake.move(direction)

        # 判断是否吃到食物
        if snake.body[0].colliderect(food.rect):
            food.spawn()
            snake.grow()
            score += 10

        # 绘制背景和得分
        screen.fill(BLACK)
        draw_score(screen, score)

        # 绘制贪吃蛇和食物
        snake.draw(screen)
        food.draw(screen)

        # 判断是否撞墙或者撞到自己
        if snake.hit_wall() or snake.hit_self():
            game_over = True

        # 更新屏幕
        pygame.display.update()

        # 设置游戏帧率
        clock.tick(FPS)

    # 游戏结束时显示分数
    screen.fill(BLACK)
    draw_text(screen, "Game Over!", SCREEN_WIDTH // 2 - 80, SCREEN_HEIGHT // 2 - 30)
    draw_text(screen, "Score: " + str(score), SCREEN_WIDTH // 2 - 60, SCREEN_HEIGHT // 2 + 10)
    pygame.display.update()

    # 等待3秒后自动退出游戏
    pygame.time.wait(3000)
    pygame.quit()
    sys.exit()


if __name__ == '__main__':
    main()
